#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <unistd.h>
#include <math.h>
//#include <windows.h> //WORD, DWORD, LONG

//constants taken from CS4953 bitmap c++
#define BI_RGB        0L        // (no compression)
#define BI_RLE8       1L        // (run-length encoding, 8 bits)
#define BI_RLE4       2L        // (run-length encoding, 4 bits)
#define BI_BITFIELDS  3L
#define BI_JPEG       4L
#define BI_PNG        5L

#define KEYBITSIZE 11
#define MAXSZ (1 << sizeof(int)-KEYBITSIZE)-1

//windows.h not available, had to make new typdefs
typedef unsigned int DWORD;
typedef unsigned short WORD;
typedef int LONG;
typedef unsigned char BYTE;

//###the following structure information is taken from wingdi.h###
//Bitmap File Header
typedef struct tagBITMAPFILEHEADER{ //14 bytes
   WORD    bfType;
   DWORD   bfSize;      // total length of bitmap file
   DWORD   sizeKEY;      //contains the message size in the first 21 bits
                        //and contains the key in the last 11 bits
   DWORD   bfOffBits;   // offset to start of actual pixel data
} BITMAPFILEHEADER;
//DIB Header
typedef struct tagBITMAPINFOHEADER{ //40 bytes
   DWORD   biSize;              // size of BITMAPINFOHEADER
   LONG    biWidth;             // width in pixels
   LONG    biHeight;            // height in pixels
   WORD    biPlanes;            // always 1
   WORD    biBitCount;          // color bits per pixel
   DWORD   biCompression;       // BI_RGB, BI_RLE8, BI_RLE4
   DWORD   biSizeImage;         // total bytes in image
   LONG    biXPelsPerMeter;     // 0, or optional horizontal resolution
   LONG    biYPelsPerMeter;     // 0, or optional vertical resolution
   DWORD   biClrUsed;           // colors actually used (normally zero, can be lower than biBitCount)
   DWORD   biClrImportant;      // important colors actualy used (normally zero)
} BITMAPINFOHEADER;
typedef struct tagRGBQUAD{
   BYTE    rgbBlue;
   BYTE    rgbGreen;
   BYTE    rgbRed;
} RGBQUAD;

static void displayFileInfo(char *pFileName, BITMAPFILEHEADER *pFileHdr,
                            BITMAPINFOHEADER *pFileInfo, RGBQUAD *pColorTable, unsigned char *pData);
static void readBMP(char *file, unsigned int *fileSize, unsigned char **image);
static unsigned int writeBMP(unsigned char *image, unsigned int fileSize, char *fileName);
static unsigned int writeExtracted(unsigned char *message, unsigned int fileSize, char *fileName);
static unsigned int hidedata(unsigned char *data, unsigned char *message, unsigned char key, unsigned int numMessages, unsigned long long numPixels);
static BYTE *extractdata(BITMAPFILEHEADER *fileHeader, BYTE *data);
static BYTE decideNumBytes(BYTE value, int numbits);
static BYTE getBits(BYTE value, int n);


/*ATTN DARREN: to set messageSz and Key
   key = (unsigned char)(bmpHeader->sizeKEY % (unsigned int)pow(2,KEYBITSIZE));
   messageSz = (unsigned int)bmpHeader->sizeKEY >> KEYBITSIZE;
*/
DWORD msgSz;
LONG imgLen;


int main(int argc, char *argv[]){
   unsigned char *image, *data, *message, flag;
   unsigned int hiddenSz, fileSz, messageSz;
   unsigned short seedVal;
   unsigned long long numPix;
   BITMAPFILEHEADER *bmpHeader;
   BITMAPINFOHEADER *bmpInfoHeader;
   RGBQUAD *bmpColorTable;

   //as we increase complexity, make sure to error check argument list passed.
   //IE, update "usage:"to
   if(argc < 3){
      fprintf(stderr,"Usage: hideex <COVER.bmp> <HIDDENDATA.*> <-x | seed INT>\n");
      fprintf(stderr,"The flag -x will make the program attempt to extract a hidden message from the file COVER.\n");
      fprintf(stderr,"If this flag is not present, the program attempts to hide whatever data is given by HIDDENDATA within COVER.\n");
      exit(1);
   }
   if(strcmp(argv[3], "-x") == 0){
      flag = 1; //flag indicates if extracting or embedding
   }else{
      seedVal = atoi(argv[3]);
      //the random key MUST be smaller than 11 bits
      if(seedVal > pow(2,11)){
         fprintf(stderr,"seed is to large, value must fit in 11 bits\n");
         exit(1);
      }
      flag = 0;
   }


   //read in source *.bmp
   readBMP(argv[1], &fileSz, &image);
   
   //begin reading in bitmap headers here
   bmpHeader = (BITMAPFILEHEADER *) (image - 2);
   bmpInfoHeader = (BITMAPINFOHEADER *) (image + 14);
   bmpColorTable = (RGBQUAD *) (image + 14 + bmpInfoHeader->biSize);
   data = image + bmpHeader->bfOffBits;

   imgLen = bmpInfoHeader->biWidth;
   
   //This ensures that the file passed in is a bitmap.
   unsigned char *indexOfFileType = (unsigned char *) strstr(image, "BM");
   if(indexOfFileType != image){
      fprintf(stderr, "%s is not a Windows Bitmap.\n", argv[1]);
   }else if(bmpInfoHeader->biBitCount != 24){
      fprintf(stderr, "%s is not 24-bit bmp\n", argv[1]);
   }
   else
      displayFileInfo(argv[1], bmpHeader, bmpInfoHeader, bmpColorTable, data);

   //begin extract/embed
   if(flag){//extracting
      fprintf(stderr, "extracting...\n");
      message = extractdata(bmpHeader, data);
      //TODO: Use correct size here, given by bmpHeader. Also, hide ALL of the message.
      writeExtracted(message, bmpHeader->sizeKEY >> KEYBITSIZE, argv[2]);
      free(message);
   }
   else{//embedding
      //readBMP for reading in the message;  WILL ACCOMPLISH THE SAME THING!
      readBMP(argv[2],&messageSz, &message);
      printf("Message: %s\tSize: %d\n", message, messageSz);
      msgSz = messageSz;

      numPix = bmpInfoHeader->biWidth * bmpInfoHeader->biHeight;
      //embed data into cover
      hiddenSz = hidedata(data, message, seedVal, messageSz, numPix);
      printf("hiddenSz:%d\n",hiddenSz);

      //insert the message size and key into the FILEHEADER, sizeKey
      bmpHeader->sizeKEY = hiddenSz;
      bmpHeader->sizeKEY = bmpHeader->sizeKEY << KEYBITSIZE;
      bmpHeader->sizeKEY = bmpHeader->sizeKEY | seedVal;

      //write out data
      writeBMP(image, fileSz, argv[1]);
   }

   return 0;
}

static void displayFileInfo(char *pFileName, BITMAPFILEHEADER *pFileHdr,
			BITMAPINFOHEADER *pFileInfo, RGBQUAD *pColorTable, unsigned char *pData)
{
   
   int numColors;

   printf("\nFile Info for %s: \n\n", pFileName);

   printf("File Type: BM\n");
   printf("File Size:%d\nData Offset:%x\n\n", pFileHdr->bfSize, (pFileHdr->bfOffBits));


   numColors = 1 << (pFileInfo->biBitCount);
   printf("Bit Map Image Info:\n\nSize:%d\nWidth:%d\nHeight:%d\nPlanes:%d\n"

        "Bits/Pixel:%d ==> %d colors\n"

        "Compression:%d\nImage Size:%d\nRes X:%d\nRes Y:%d\nColors:%d\nImportant Colors:%d\n\n",

        pFileInfo->biSize,

        pFileInfo->biWidth,

        pFileInfo->biHeight,

        pFileInfo->biPlanes,

        pFileInfo->biBitCount, numColors,

        pFileInfo->biCompression,

        pFileInfo->biSizeImage,

        pFileInfo->biXPelsPerMeter,

        pFileInfo->biYPelsPerMeter,

        pFileInfo->biClrUsed,

        pFileInfo->biClrImportant);
}

static unsigned int writeBMP(unsigned char *image, unsigned int fileSize, char *fileName){
   FILE *outfile;
   char fileBuffer[256];
   int offset;
   unsigned int bytesWritten;

   //add .hx.bmp to stego cover
   offset = strlen(fileName)-4;
   strncpy(fileBuffer, fileName, offset);
   fileBuffer[offset] = '\0';
   strcat(fileBuffer, ".hx.bmp");

   //begin write
   if((outfile = fopen(fileBuffer, "w+"))==NULL){
      fprintf(stderr,"error writing to %s\n",fileBuffer);
      perror(NULL);
      exit(1);
   }
   if((bytesWritten = fwrite(image, sizeof(unsigned char), fileSize, outfile)) != fileSize){
      fprintf(stderr, "incorrect write to file. bytes to write: %d, bytes written: %d\n", fileSize, bytesWritten);
      perror("incorrect write to file\n");
      exit(1);
   }
   fclose(outfile);

   return bytesWritten;
}
static unsigned int writeExtracted(unsigned char *message, unsigned int fileSize, char *fileName){
   FILE *outfile;
   char fileBuffer[256];
   unsigned int bytesWritten;

   //begin write
   if((outfile = fopen(fileName, "w+"))==NULL){
      fprintf(stderr,"error writing to %s\n",fileBuffer);
      perror(NULL);
      exit(1);
   }
   if((bytesWritten = fwrite(message, sizeof(unsigned char), fileSize, outfile)) != fileSize){
      fprintf(stderr, "incorrect write to file. bytes to write: %d, bytes written: %d\n", fileSize, bytesWritten);
      perror("incorrect write to file\n");
      exit(1);
   }
   fclose(outfile);

   return bytesWritten;
}

static void readBMP(char *file, unsigned int *fileSize, unsigned char **image){
   FILE *infile;

   if((infile = fopen(file,"r")) == NULL){
      fprintf(stderr,"error opening file %s\n",file);
      perror(NULL);
      exit(1);
   }

   fseek(infile,0,SEEK_END);
   *fileSize = ftell(infile);
   fseek(infile,0,SEEK_SET);

   if((*image = (unsigned char *) malloc(*fileSize)) == NULL){
      perror("malloc error");
      exit(1);
   }
   fread(*image, sizeof(unsigned char), *fileSize, infile);
   fclose(infile);
}

static unsigned int hidedata(unsigned char *data, unsigned char *message, unsigned char key, unsigned int numMessages, unsigned long long numPixels){
   int numbits, n, h;
   unsigned int m, totalHidden;
   unsigned long long p;
   WORD bitbuf;
   DWORD pixelCounter;
   RGBQUAD *pixelBuf =(RGBQUAD *) data; // represents current pixel, assigning the first pixel from bmp
   BYTE *temp; // Used to offset RGBQUAD pixel if hit with a end-of-row buffer
   BYTE hidebuf; // represents current message
   BYTE bitOffset; // used to replace n bits to 0
   BYTE diff; // Used for row buffers; Pixels Per Row * 3 Bytes % 4 Bytes Req. Per Row
   BYTE r, g, b; // Holds RGB value to use in condtions and calculations
   h = 0; // number of bits hidden
   m = 0; // current index of info
   p = 0; // current index of pixel
   bitbuf = 0; // bits of data that are ready to be hidden
   numbits = 0; // number of bits in bitbuf
   pixelCounter = 1; // Pixels read in per row
   numMessages *= 8;
   totalHidden = 0; //TODO DEBUG

   int z; BYTE y;
   // Set and Hide Key
   srand(key); // seed rand to key

   do {
      // Read in Hidden Data -- Combine next byte of info into bitbuf
      // Note: One byte can only be read in, and 3~5 bits may be hidden
      //printf("NumBits:%d\tNumMsgs:%d\n",numbits,numMessages);
      if(numbits < 5 && (m+1)*8 < numMessages) {
         hidebuf = (BYTE) *(message + m); // gets next byte from the message
         bitbuf = (bitbuf << 8) + hidebuf; // combines bits into a single stream
         numbits += 8;
         m++;
         //printf("Read: %d\tBuffer: %d\tSize:%d\n",hidebuf, bitbuf,numbits);
      }

      // assign to shorten code
      r = pixelBuf->rgbRed;
      g = pixelBuf->rgbGreen;
      b = pixelBuf->rgbBlue;

      // Choose indicator channel (case) -- 0,1,2 is R,G,B respectively
      // Replace LSB: leftmost channel is 0, and rightmost is 1
      switch(z = rand() % 3) {
         //printf("Indicator:%d\t", z);
         case 0: // RED
            if(g <= b) { // decide which of the 2 remaining channels is lower
               // write RED + indicator
               pixelBuf->rgbRed = r & 254; // LSB=0 (left channel is indicator)
               // write GREEN + hidden
               n = decideNumBytes(g, numbits);
               bitOffset = 256 - pow(2,n); // Ex. n=3 11111111-0000111
               y = pixelBuf->rgbGreen = (g & bitOffset) + (bitbuf >> (numbits-n));
            }
            else {
               // write RED + indicator
               pixelBuf->rgbRed = r | 1; // LSB=1 (right channel is indicator)
               // write BLUE + hidden
               n = decideNumBytes(b, numbits);
               bitOffset = 256 - pow(2,n);
               y = pixelBuf->rgbBlue = (b & bitOffset) + (bitbuf >> (numbits-n));
            }
            break;
         case 1: // GREEN
            if(r < b) {
               // write RED + hidden
               n = decideNumBytes(r, numbits);
               bitOffset = 256 - pow(2,n);
               y = pixelBuf->rgbRed = (r & bitOffset) + (bitbuf >> (numbits-n));
               // write GREEN + indicator
               pixelBuf->rgbGreen = g & 254; // LSB is 0
            }
            else {
               // write GREEN + indicator
               pixelBuf->rgbGreen = g | 1; // LSB is 1
               // write BLUE + hidden
               n = decideNumBytes(b, numbits);
               bitOffset = 256 - pow(2,n);
               y = pixelBuf->rgbBlue = (b & bitOffset) + (bitbuf >> (numbits-n));
            }
            break;
         case 2: // BLUE
            if(r <= g) {
               // write RED + hidden
               n = decideNumBytes(r, numbits);
               bitOffset = 256 - pow(2,n);
               y = pixelBuf->rgbRed = (r & bitOffset) + (bitbuf >> (numbits-n));
               // write BLUE + indicator
               pixelBuf->rgbBlue = b & 254; // LSB is 0
            }
            else {
               // write GREEN + hidden
               n = decideNumBytes(g, numbits);
               bitOffset = 256 - pow(2,n);
               y = pixelBuf->rgbGreen = (g & bitOffset) + (bitbuf >> (numbits-n));
               // write BLUE + indicator
               pixelBuf->rgbBlue = b | 1; // LSB is 1
            }
      } // end switch
      printf("R%d\tG%d\tB%d\tIndicator: %d\n",r,g,b,z);
      printf("\tN: %d\tReplace With: %d\tNew Value: %d\n", n,(bitbuf >> (numbits-n)),y);
      // Update bitbuf
      bitbuf = bitbuf << (16-numbits+n); // Separated -- if on a single line
      bitbuf = bitbuf >> (16-numbits+n); //    even with ()'s, cancels out
      numbits -= n;
      totalHidden += n;

      // End of row offset
      if((pixelCounter*3)%4 != 0 && p % imgLen == 0) {
         diff = 4-((pixelCounter)*3)%4;
         temp = (BYTE*)pixelBuf; // To offset a RGBQUAD (3 BYTES) by diff amount of byte
         temp += diff;
         pixelBuf = (RGBQUAD*)temp; // Convert back;
         pixelCounter = 0; // Reset counter for next row
      }

      printf("m:%d\tnumMsgs:%d\tn:%d\n",(m+1)*8,numMessages,n);
      printf("p = %llu, numPixels = %llu\n\n",p,numPixels);

      p++;
      pixelBuf++;
      pixelCounter++;
      //printf("Indicator:%d\tChanged:%d\n",z,y);
      //printf("Update Buffer: %d\tSize:%d\n",bitbuf,numbits);

      printf("totalHidden:%d\treturns:%d\n",totalHidden,m);
   } // end while
   while(p < numPixels && ((m+1)*8 < numMessages || numbits > 0));

   if((m+1)*8 < numMessages)
   printf("NOTE: Not enough capacity to hold all of the hidden message!\n");

   //return m; // Returns Number of BYTES
   return totalHidden; // Returns number of BITS hidden
} // end hide data

// Assume *data [start of first pixel] is global
static BYTE *extractdata(BITMAPFILEHEADER *fileHeader, BYTE *data) {
   int numbits, m, n, p;
   DWORD messageSz, param, pixelCounter;
   WORD bitbuf, h, maxSize;
   RGBQUAD *pixelBuf = (RGBQUAD*) data; // represents current pixel
   BYTE *temp; // Used to offset RGBQUAD pixel if hit with a end-of-row buffer
   BYTE *extracted; // represents current extracted message
   BYTE r, g, b; // Holds RGB value to use in condtions and calculations
   BYTE diff; // Used for row buffers; Pixels Per Row * 3 Bytes % 4 Bytes Req. Per Row
   BYTE key;
   h = 0; // total number of hidden bits extracted
   m = 0; // current index of info
   p = 0; // current index of pixel
   bitbuf = 0; // bits of data that are ready to be hidden
   numbits = 0; // number of bits in bitbuf
   pixelCounter = 1; // Pixels read in per row

   // extract paramaters from header
   param = fileHeader->sizeKEY;
   key = (BYTE)(fileHeader->sizeKEY % (DWORD)pow(2,KEYBITSIZE));
   messageSz = (DWORD)fileHeader->sizeKEY >> KEYBITSIZE;

   // Set Key
   //extracted = malloc(sizeof(BYTE)*messageSz); // OLD
   maxSize = ceil(messageSz/8.0);
   extracted = malloc(sizeof(BYTE)*maxSize);
   srand(key);

   printf("maxSize in Bytes:%d\n",maxSize);

   int z; // For Debug -- Indicator channel
   do {
      printf("m:%d begins\tmaxSizeBYTES:%d\t",m,maxSize);
      // Extract Hidden Data -- Combine next byte of info into bitbuf
      // Note: One byte can only be read in, and 3~5 bits may be hidden
      if(numbits > 7) {
         extracted[m] = bitbuf >> (numbits-8); // Store byte into memory, then update bitbuf
         bitbuf = bitbuf << (16-numbits+8); // Separated -- if on a single line
         bitbuf = bitbuf >> (16-numbits+8); //    even with ()'s, cancels out
         numbits -= 8;
         //printf("Stored: %d\tBuffer: %d\tSize:%d\n",extracted[m], bitbuf,numbits);
         m++;
      }

      // assign to shorten code
      r = pixelBuf->rgbRed;
      g = pixelBuf->rgbGreen;
      b = pixelBuf->rgbBlue;

      // Find indicator channel (case) -- 0,1,2 is R,G,B respectively
      // Replace LSB: leftmost channel is 0, and rightmost is 1
      switch(z = rand() % 3) {
         case 0: // RED
            if(r % 2 == 0) { // LSB is 0 (Green is hidden channel)
               n = decideNumBytes(g, messageSz-h);
               bitbuf = (bitbuf << n) + getBits(g, n);
            }
            else { // LSB is 1 (Blue is hidden channel)
               n = decideNumBytes(b, messageSz-h);
               bitbuf = (bitbuf << n) + getBits(b, n);
            }
            break;
         case 1: // GREEN
            if(g % 2 == 0) { // LSB is 0 (Red is hidden channel)
               n = decideNumBytes(r, messageSz-h);
               bitbuf = (bitbuf << n) + getBits(r, n);
            }
            else { // LSB is 1 (Blue is hidden channel)
               n = decideNumBytes(b, messageSz-h);
               bitbuf = (bitbuf << n) + getBits(b, n);
            }
            break;
         case 2: // BLUE
            if(b % 2 == 0) { // LSB is 0 (Red is hidden channel)
               n = decideNumBytes(r, messageSz-h);
               bitbuf = (bitbuf << n) + getBits(r, n);
            }
            else { // LSB is 1 (Green is hidden channel)
               n = decideNumBytes(g, messageSz-h);
               bitbuf = (bitbuf << n) + getBits(g, n);
            }
      } // end switch
      //printf("R:%d\tG:%d\tB:%d\tIndicator:%d\n",r,g,b,z);
      //printf("Pixel:%d\tN: %d\tNew Buffer: %d\n",p,z, n,bitbuf);
      numbits += n;
      h += n;

      if((pixelCounter*3)%4 != 0 && p % imgLen == 0) {
         diff = 4-((pixelCounter)*3)%4;
         temp = (BYTE*)pixelBuf; // To offset a RGBQUAD (3 BYTES) by diff amount of byte
         temp += diff;
         pixelBuf = (RGBQUAD*)temp; // Convert back;
         pixelCounter = 0; // Reset counter for next row
      }

      p++;
      pixelBuf++;
      pixelCounter++;

      //if(m+1 == messageSz && numbits > 1) { // all hidden bits extracted OLD
      if(h >= messageSz && numbits > 1) {
         extracted[m] = bitbuf << (8-numbits);
         //printf("Stored: %d\tBuffer: %d\tSize:%d\n",extracted[m], bitbuf,numbits);
         //m++;
      }

      //printf("totalExtracted:%d\tm:%d\n\n",h,m);
      printf("m:%d ends\n",m);
   } // end while
   //while(m <= messageSz); // old
   while(h < messageSz && m < maxSize);

   printf("messageSz:%d\tmaxSize:%d\n",messageSz,maxSize);

   int i;
   //for(i=0; i<messageSz; i++) // old
   for(i=0; i<maxSize; i++)
      fprintf(stderr, "%c",extracted[i]);
   fprintf(stderr,"\n");
   return extracted;
} // end extract data

// Color's value determines how many bits to replace
// Note: If numbits is less than 4, then the number is equal to numbits
static BYTE decideNumBytes(BYTE value, int numbits){
   // Note: Thresholds were decided through visual perception trial & error
   if(numbits < 4)
      return numbits;
   if(value < 32)
      return 5;
   if(value < 64)
      return 4;
   return 3;
}

// Returns n number of hidden bits from the given value
static BYTE getBits(BYTE value, int n) {
   BYTE offset = value << (8-n); // Separated -- if on a single line
   return offset >> (8-n);       //    even with ()'s, cancels out
}
